import numpy as np
import csv
from matplotlib import pyplot as plt

ACC_LIMIT = 3 # the limit of acceleration, this can be calibrated based on the data
Ts = 0.04 # time interval for data sampling for HighD is 0.04 for other datasets are 0.1
# Define IDM model
def IDM(para, spacing, svSpd, relSpd):

    """
    Funciton that takes IDM paramters and car-following states as inputs, and output
    the acceleration for the following vehicle.

    :param para: a vector containing IDM's parameters.
        E.g., para = np.array([32.0489077 ,  0.74084102,  1.18623382,  0.87773747,  1.,2.95210611])
    :param spacing: scaler, gap between two vehicles [m].
    :param svSpd: speed of the following vehicle [m/s].
    :param relSpd: lead vehicle speed - following vehicle speed [m/s].
    :return: acceleration of the following vehicle in next step [m/s^2].
    """
    desiredSpd = para[0]  # in m/s
    desiredTimeHdw = para[1]  # in seconds
    maxAcc = para[2]  # m/s^2
    comfortAcc = para[3]  # m/s^2
    beta = para[4]
    jamSpace = para[5]  # in meters

    desiredSpacing = jamSpace + max(0, desiredTimeHdw * svSpd - svSpd * relSpd / (2 * np.sqrt(maxAcc * comfortAcc)))
    acc = maxAcc * (1 - (svSpd / desiredSpd) ** beta - (desiredSpacing / spacing) ** 2)

    return acc
def simulate_car_fol_IDM_test(model_fun, lvSpd, init_s, init_svSpd, para):
    """
    Simulate a car following event based on a car-following model.

    :param model_fun:
    """
    T = Ts  # data sampling interval

    svSpd_sim = []
    spacing_sim = []
    # IDM
    spacing, svSpd, relSpd = init_s, init_svSpd, lvSpd[0] - init_svSpd
    svSpd_sim.append(svSpd)
    spacing_sim.append(spacing)

    is_collision=False
    for i in range(1, len(lvSpd)):
        # calcualate next_step acceleration using IDM model
        acc = model_fun(para, spacing, svSpd, relSpd)

        # state update based on Newton's motion law
        svSpd_ = max(0.001, svSpd + acc * T)  # next step svSpd
        relSpd_ = lvSpd[i] - svSpd_
        spacing_ = spacing + T * (relSpd_ + relSpd) / 2

        # update state variables
        svSpd = svSpd_
        relSpd = relSpd_
        spacing = spacing_
        svSpd_sim.append(svSpd)
        spacing_sim.append(spacing)

    return np.asarray(svSpd_sim), np.asarray(spacing_sim)

#在文件处上传要使用的轨迹文件
filename = 'new.csv'  # 要读取的CSV文件路径
#这里提取lv的速度值
column_index = 2  # 要提取的列的索引（从0开始）
start_row = 118298  # 要提取的起始行索引（从0开始）
end_row = 118748  # 要提取的结束行索引（包含在内）

# 打开CSV文件
with open(filename, 'r') as file:
    reader = csv.reader(file)

    # 跳过起始行之前的行
    for _ in range(start_row):
        next(reader)

    # 逐行读取CSV文件并提取指定列的数据
    data = []
    for row in reader:
        data.append(row[column_index])

        # 判断是否达到结束行
        if reader.line_num == end_row:
            break
#提取得到的lv
lv=np.array(data,dtype = 'float64')

# Simulate IDM with calibrated para
#通过计算数据集，得到的参数值
para =[31.63398677,  0.8485575 ,  0.35958698 , 0.39194875 , 2.89009438 , 3.4968872 ] # HighD
#initial_space,预测时，必须指定一个初始的与后车的距离
#initial_speed,预测时，必须指定一个初始的后车速度
initial_space=50
initial_speed=22
#speed_sim是预测的后车速度，spacing_sim是预测的前后两车的距离
speed_sim, spacing_sim = simulate_car_fol_IDM_test(IDM, lv, initial_space, initial_speed, para)
plt.figure()
plt.plot(lv, label='LV',c='orange')
plt.plot(speed_sim,'--', label='SV_sim',c='deepskyblue')
plt.legend()
plt.xlabel('Time step (0.04s as sample interval)')
plt.ylabel('Speed (m/s)')
plt.figure()
plt.plot(spacing_sim,'--', label='spacing_sim ',c='deepskyblue')
plt.legend()
plt.xlabel('Time step (0.04s as sample interval)')
plt.ylabel('Spacing (m)')
plt.show()